from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import r2_score

try:
    from sklearn.cross_validation import KFold
    __USE_OLD_KOLD = True
except Exception:
    from sklearn.model_selection import KFold
    __USE_OLD_KOLD = False


import numpy as np
import torch
import torch.nn
from torch import nn
from torch.autograd import Variable

def _T(obj):
    return Variable(torch.from_numpy(np.array(obj, dtype=np.float32)))

def make_split(X, y):
    ri = np.random.permutation(X.shape[0])
    X = X.copy()[ri]
    y = y.copy()[ri]
    length = X.shape[0]
    sp1 = int(length * 0.8)
    sp2 = int(length * 0.9)
    return X[:sp1, :], X[sp1:sp2, :], X[:sp2], X[sp2:], y[:sp1], y[sp1:sp2], y[:sp2], y[sp2:]



def get_train_test(X, y, kfold=False, kfold_k=None, kfold_shuffle=False):
    if kfold is False:
        yield X, y, None, None
    else:
        if __USE_OLD_KOLD:
            kf = KFold(X.shape[0], n_folds=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf:
                yield X[train_index], y[train_index], X[test_index], y[test_index]
        else:
            kf = KFold(n_splits=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf.split(X):
                yield X[train_index], y[train_index], X[test_index], y[test_index]


def get_train_test_index(X, y, kfold_k=4, kfold_shuffle=True):
    if __USE_OLD_KOLD:
        kf = KFold(X.shape[0], n_folds=kfold_k, shuffle=kfold_shuffle)
        for train_index, test_index in kf:
            yield train_index, test_index
    else:
        kf = KFold(n_splits=kfold_k, shuffle=kfold_shuffle)
        for train_index, test_index in kf.split(X):
            yield train_index, test_index

def mape_score(labels, predictions):  # 77.69%
    predictions = np.array(predictions)
    labels = np.array(labels)
    if (labels == 0).any():
        print('warning: {} labels = 0, will be dropped.'.format((labels == 0).sum()))
    predictions = predictions[labels != 0]
    labels = labels[labels != 0]
    return np.abs((predictions - labels) / (labels)).mean()


def kfold_validation(X, y, kfold_k, model, vis=False, verbose=False, use_mape=False):
    all_score = 0
    if use_mape:
        prefix = "1 - mape score"
    else:
        prefix = "model score"
    for i, (X_train, y_train, X_test, y_test) in enumerate(get_train_test(X, y, True, kfold_k, True)):
        clf = model
        clf.fit(X_train, y_train)
        if use_mape:
            score = 1 - mape_score(y_test, clf.predict(X_test))
        else:
            score = clf.score(X_test, y_test)
        if verbose:
            print('%s: %d / %d: %f' % (prefix, i + 1, kfold_k, score))
        all_score += score
        if vis:
            y_predicted = clf.predict(X_test)
            fig, ax = plt.subplots()
            ax.scatter(y_test, y_predicted)
            ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=4)
            ax.set_xlabel('Observed')
            ax.set_ylabel('Predicted')
            plt.show()
    if verbose:
        print('avg score: %f' % (all_score / kfold_k))
    return all_score / kfold_k


def validation_tf(model, dataset, batch_size, max_step, numpy_dataset_test, verbose=True):
    tf.reset_default_graph()

    opt = tf.train.GradientDescentOptimizer(learning_rate=0.01)

    cal_value_gradients = tfe.implicit_value_and_gradients(model.loss)

    batched_dataset = dataset.batch(batch_size)

    r2s = []
    oneminus_mapes = []
    steps = []

    print('use_feature: feature_ab')
    for (batch, input_batch) in enumerate(tfe.Iterator(batched_dataset)):
        if batch + 1 > max_step:
            break
        loss, grads = cal_value_gradients(input_batch, input_batch['target'])
        opt.apply_gradients(grads)
        if (batch + 1) % 100 == 0:
            r2 = model.score(numpy_dataset_test, numpy_dataset_test['target'], metrics="r2")
            mape = model.score(numpy_dataset_test, numpy_dataset_test['target'], metrics="mape")
            if verbose:
                print('step: {} /  {}  loss: {}  r2_score: {}  1 - mape_score: {}'.format(batch + 1, max_step, loss.numpy(), r2, 1 - mape))
            r2s.append(r2)
            oneminus_mapes.append(1 - mape)
            steps.append(batch + 1)
    r2 = model.score(numpy_dataset_test, numpy_dataset_test['target'], metrics="r2")
    mape = model.score(numpy_dataset_test, numpy_dataset_test['target'], metrics="mape")
    r2s.append(r2)
    oneminus_mapes.append(1 - mape)
    steps.append(batch)
    if verbose:
        print('final step: {}  r2_score: {}  1 - mape_score: {}'.format(batch, r2, 1 - mape))
    return np.array(steps), np.array(r2s), np.array(oneminus_mapes)

